import torch
import numpy as np
import pandas as pd
import torch.nn.functional as F

# 1. 数据准备
df_train_data = pd.read_csv('abalone_train.data')
df_test_data = pd.read_csv('abalone_test.data')

train_data_x = df_train_data.iloc[:, 1:10]
train_data_x = np.array(train_data_x, dtype=float)
train_data_x = torch.squeeze(torch.FloatTensor(train_data_x), dim=1)

test_data_x = df_test_data.iloc[:, 1:10]
test_data_x = np.array(test_data_x, dtype=float)
test_data_x = torch.squeeze(torch.FloatTensor(test_data_x), dim=1)

df_train_data['Sex Code'] = df_train_data['Sex'].map({'M': 0, 'F': 1, 'I': 2})
df_test_data['Sex Code'] = df_test_data['Sex'].map({'M': 0, 'F': 1, 'I': 2})

train_data_y = df_train_data['Sex Code']
train_data_y = np.array(train_data_y, dtype=int)
train_data_y = torch.LongTensor(train_data_y)

test_data_y = df_test_data['Sex Code']
test_data_y = np.array(test_data_y, dtype=int)
test_data_y = torch.LongTensor(test_data_y)

input_data = train_data_x
label = train_data_y

if torch.cuda.is_available():
    input_data = input_data.cuda()
    label = label.cuda()
    test_data_x = test_data_x.cuda()


# 2. 定义BP神经网络
class Net(torch.nn.Module):
    def __init__(self, n_feature, n_hidden, n_output, dropout=0.5):
        super(Net, self).__init__()

        self.hidden = torch.nn.Linear(n_feature, n_hidden)  # hidden layer
        self.out = torch.nn.Linear(n_hidden, n_output)  # 定义输出层网络

    def forward(self, x):
        x = F.relu(self.hidden(x))  # 隐藏层的激活函数，采用relu，也可以采用sigmoid, tanh
        x = self.out(x)  # 输出层不用激活函数
        return x


# 3. 定义优化器损失函数
net = Net(n_feature=8, n_hidden=200, n_output=3)  # n_feature:输入的特征维度, n_hidden:神经元个数, n_output:输出的类别个数
if torch.cuda.is_available():
    net = net.cuda()
optimizer = torch.optim.Adam(net.parameters(), lr=0.015)
loss_func = torch.nn.CrossEntropyLoss()  # 对于多分类一般采用的交叉熵损失函数
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20000, gamma=0.9, last_epoch=-1)

# 4. 训练数据
for epoch in range(100001):
    out = net(input_data)  # 输入input,输出out
    loss = loss_func(out, label)  # 输出与label对比
    if epoch % 1000 == 0:
        print("epoch: {}, loss: {:.6f}".format(epoch, loss))
    optimizer.zero_grad()  # 梯度清零
    loss.backward()  # 前馈操作
    optimizer.step()  # 使用梯度优化器
    scheduler.step()

# 5. 得出结果
if torch.cuda.is_available():
    torch.cuda.empty_cache()

out = net(test_data_x)  # out是一个计算矩阵，可以用Fun.softmax(out)转化为概率矩阵
prediction = torch.max(out, 1)[1]  # 返回index  0返回原值
if torch.cuda.is_available():
    # 将GPU的数据转存到CPU
    prediction = prediction.cpu()
pred_y = prediction.data.numpy()
target_y = test_data_y.data.numpy()

# 6. 衡量准确率
accuracy = float((pred_y == target_y).astype(int).sum()) / float(target_y.size)
print("=> prediction accuracy: {:.6f}".format(accuracy))
